#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
实时协作研究演示脚本
演示如何使用WebSocket进行多用户实时协作研究
"""

import asyncio
import json
import websockets
import threading
import time
from typing import Dict, Any

class CollaborativeResearchDemo:
    """实时协作研究演示类"""
    
    def __init__(self):
        self.ws_url = "ws://localhost:8000/ws/collaborate"
        self.session_id = "demo_session_1"
        self.users = []
        
    async def simulate_user(self, user_id: str, actions: list):
        """
        模拟用户行为
        
        Args:
            user_id: 用户ID
            actions: 用户要执行的操作列表
        """
        uri = f"{self.ws_url}?session_id={self.session_id}&user_id={user_id}"
        
        try:
            async with websockets.connect(uri) as websocket:
                print(f"用户 {user_id} 连接到会话 {self.session_id}")
                
                # 接收连接确认消息
                message = await websocket.recv()
                print(f"用户 {user_id} 收到消息: {message}")
                
                # 执行用户操作
                for action in actions:
                    if action["type"] == "delay":
                        await asyncio.sleep(action["duration"])
                    elif action["type"] == "send":
                        await websocket.send(json.dumps(action["message"]))
                        print(f"用户 {user_id} 发送消息: {action['message']}")
                    elif action["type"] == "listen":
                        try:
                            # 设置接收超时
                            response = await asyncio.wait_for(websocket.recv(), timeout=action.get("timeout", 5))
                            print(f"用户 {user_id} 收到响应: {response}")
                        except asyncio.TimeoutError:
                            print(f"用户 {user_id} 等待响应超时")
                
                # 保持连接一段时间
                await asyncio.sleep(2)
                
        except Exception as e:
            print(f"用户 {user_id} 连接错误: {e}")
    
    def run_demo(self):
        """运行演示"""
        print("=== 实时协作研究演示 ===")
        print("启动多用户实时协作研究演示...")
        
        # 定义用户行为
        user1_actions = [
            {
                "type": "send",
                "message": {
                    "type": "chat_message",
                    "content": "大家好！我们开始研究'人工智能在医疗领域的应用'这个话题吧",
                    "timestamp": time.time()
                }
            },
            {
                "type": "delay",
                "duration": 1
            },
            {
                "type": "send",
                "message": {
                    "type": "research_start",
                    "question": "人工智能在医疗领域的应用有哪些最新进展？",
                    "timestamp": time.time()
                }
            },
            {
                "type": "listen",
                "timeout": 10
            }
        ]
        
        user2_actions = [
            {
                "type": "delay",
                "duration": 2
            },
            {
                "type": "send",
                "message": {
                    "type": "chat_message",
                    "content": "我来关注诊断方面的应用",
                    "timestamp": time.time()
                }
            },
            {
                "type": "delay",
                "duration": 1
            },
            {
                "type": "send",
                "message": {
                    "type": "cursor_position",
                    "position": {"x": 100, "y": 200, "document": "medical_ai_research.md"},
                    "timestamp": time.time()
                }
            },
            {
                "type": "listen",
                "timeout": 10
            }
        ]
        
        user3_actions = [
            {
                "type": "delay",
                "duration": 3
            },
            {
                "type": "send",
                "message": {
                    "type": "chat_message",
                    "content": "我会关注治疗方面的应用",
                    "timestamp": time.time()
                }
            },
            {
                "type": "listen",
                "timeout": 10
            }
        ]
        
        # 创建异步任务
        async def run_users():
            tasks = [
                asyncio.create_task(self.simulate_user("user1", user1_actions)),
                asyncio.create_task(self.simulate_user("user2", user2_actions)),
                asyncio.create_task(self.simulate_user("user3", user3_actions))
            ]
            await asyncio.gather(*tasks)
        
        # 运行演示
        asyncio.run(run_users())
        
        print("\n演示完成！")

if __name__ == "__main__":
    demo = CollaborativeResearchDemo()
    demo.run_demo()